"""
Main DLR-FEM implementation for Allen-Cahn equation.

This module implements the complete second-order DLR-MLFEM integrator 
(Algorithm 3.2) using Strang splitting for the Allen-Cahn equation.
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
from typing import Tuple, Optional, Dict, List
import time
from dataclasses import dataclass
import sys
import os

# Add parent directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from paper_2501_06145_dlrfem.fem import create_uniform_mesh, create_basis_functions, FEMOperators, InitialConditions
from paper_2501_06145_dlrfem.low_rank import create_linear_solver, create_bug_integrator


@dataclass
class DLRFEMConfig:
    """Configuration parameters for DLR-FEM simulation."""
    
    # Domain and discretization
    domain_x: Tuple[float, float] = (0.0, 1.0)
    domain_y: Tuple[float, float] = (0.0, 1.0)
    nx: int = 64
    ny: int = 64
    poly_order: int = 1
    
    # Physical parameters
    epsilon: float = 0.02  # Interface parameter
    
    # Time integration
    T_final: float = 1.0
    dt: float = 0.001
    
    # Low-rank parameters
    initial_rank: int = 10
    max_rank: int = 50
    truncation_tol: float = 1e-12
    
    # Adaptive time stepping
    adaptive: bool = False
    energy_tol: float = 1e-6
    min_dt: float = 1e-8
    
    # Output control
    save_frequency: int = 10
    plot_frequency: int = 50
    verbose: bool = True


class DLRFEMSolver:
    """
    Complete DLR-FEM solver for Allen-Cahn equation.
    
    Implements Algorithm 3.2 (Second-Order DLR-MLFEM Integrator) using
    Strang splitting: S_{τ/2}^L ∘ S_τ^N ∘ S_{τ/2}^L
    """
    
    def __init__(self, config: DLRFEMConfig):
        """
        Initialize DLR-FEM solver.
        
        Args:
            config: Configuration parameters
        """
        self.config = config
        
        # Initialize components
        self._setup_discretization()
        self._setup_operators()
        self._setup_integrators()
        
        # Storage for solution history
        self.solution_history = []
        self.time_history = []
        self.rank_history = []
        self.energy_history = []
        
        if config.verbose:
            print(f"DLR-FEM solver initialized:")
            print(f"  Domain: {config.domain_x} × {config.domain_y}")
            print(f"  Grid: {config.nx} × {config.ny}")
            print(f"  Time step: {config.dt}")
            print(f"  Final time: {config.T_final}")
            print(f"  Epsilon: {config.epsilon}")
            print(f"  Max rank: {config.max_rank}")
    
    def _setup_discretization(self):
        """Setup mesh and basis functions."""
        self.mesh = create_uniform_mesh(
            self.config.domain_x, 
            self.config.domain_y,
            self.config.nx, 
            self.config.ny
        )
        
        self.basis = create_basis_functions(self.config.poly_order)
        
        # Get coordinate meshes
        self.X, self.Y = self.mesh.get_coordinates()
    
    def _setup_operators(self):
        """Setup finite element operators."""
        self.operators = FEMOperators(self.mesh, self.basis)
        
        # Get matrices
        self.M_x, self.M_y = self.operators.get_mass_matrices()
        self.A_x, self.A_y = self.operators.get_stiffness_matrices()
        self.L_x, self.L_y = self.operators.get_linear_operators()
    
    def _setup_integrators(self):
        """Setup time integrators."""
        # Linear solver (Algorithm 3.1)
        self.linear_solver = create_linear_solver(
            self.M_x, self.M_y, self.L_x, self.L_y,
            epsilon=self.config.epsilon,
            adaptive=self.config.adaptive
        )
        
        # Nonlinear integrator (BUG)
        self.nonlinear_solver = create_bug_integrator(
            self.M_x, self.M_y,
            max_rank=self.config.max_rank
        )
    
    def set_initial_condition(self, ic_type: str = "kiss_bubble", **kwargs) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Set initial condition and decompose into low-rank form.
        
        Args:
            ic_type: Type of initial condition
            **kwargs: Additional parameters for initial condition
            
        Returns:
            (U0, S0, V0): Initial low-rank factors
        """
        if ic_type == "kiss_bubble":
            u0 = InitialConditions.kiss_bubble(
                self.X, self.Y, 
                epsilon=self.config.epsilon,
                **kwargs
            )
        elif ic_type == "random":
            u0 = InitialConditions.random_perturbation(
                self.X, self.Y,
                **kwargs
            )
        elif ic_type == "sine":
            u0 = InitialConditions.sine_wave(
                self.X, self.Y,
                **kwargs
            )
        else:
            raise ValueError(f"Unknown initial condition type: {ic_type}")
        
        # Decompose via SVD
        U, s, Vt = np.linalg.svd(u0, full_matrices=False)
        
        # Truncate to initial rank
        r = min(self.config.initial_rank, len(s))
        self.U = U[:, :r]
        self.S = np.diag(s[:r])
        self.V = Vt[:r, :].T
        
        # Store initial condition
        self.u0 = u0
        
        if self.config.verbose:
            print(f"Initial condition set: {ic_type}")
            print(f"Initial rank: {r}")
            print(f"Initial solution range: [{np.min(u0):.3f}, {np.max(u0):.3f}]")
        
        return self.U, self.S, self.V
    
    def step(self, U: np.ndarray, S: np.ndarray, V: np.ndarray, 
             dt: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Perform one step of the DLR-MLFEM integrator (Algorithm 3.2).
        
        Strang splitting: S_{τ/2}^L ∘ S_τ^N ∘ S_{τ/2}^L
        
        Args:
            U, S, V: Current low-rank factors
            dt: Time step
            
        Returns:
            Updated low-rank factors
        """
        # Step 1: First linear half-step (τ/2)
        U_1, S_1, V_1 = self.linear_solver.step(U, S, V, dt / 2)
        
        # Step 2: Nonlinear full step (τ)
        U_2, S_2, V_2 = self.nonlinear_solver.step(U_1, S_1, V_1, dt)
        
        # Step 3: Second linear half-step (τ/2)
        U_final, S_final, V_final = self.linear_solver.step(U_2, S_2, V_2, dt / 2)
        
        return U_final, S_final, V_final
    
    def solve(self, ic_type: str = "kiss_bubble", **ic_kwargs) -> Dict:
        """
        Solve the Allen-Cahn equation using DLR-FEM.
        
        Args:
            ic_type: Initial condition type
            **ic_kwargs: Initial condition parameters
            
        Returns:
            Dictionary with solution results
        """
        if self.config.verbose:
            print("Starting DLR-FEM solution...")
        
        # Set initial condition
        U, S, V = self.set_initial_condition(ic_type, **ic_kwargs)
        
        # Time stepping
        t = 0.0
        step_count = 0
        dt = self.config.dt
        
        # Initialize storage
        self.time_history = [t]
        self.rank_history = [S.shape[0]]
        
        # Store initial energy
        if hasattr(self.linear_solver, 'energy'):
            energy = self.linear_solver.energy(U, S, V)
            self.energy_history = [energy]
        
        # Store initial solution if requested
        if step_count % self.config.save_frequency == 0:
            solution = U @ S @ V.T
            self.solution_history.append(solution.copy())
        
        start_time = time.time()
        
        while t < self.config.T_final:
            # Adaptive time stepping
            if self.config.adaptive and hasattr(self.linear_solver, 'adaptive_step'):
                U_new, S_new, V_new, dt_used = self._adaptive_step(U, S, V, dt)
                dt = min(dt_used * 1.1, self.config.dt)  # Gradually increase dt
            else:
                # Fixed time stepping
                dt_actual = min(dt, self.config.T_final - t)
                U_new, S_new, V_new = self.step(U, S, V, dt_actual)
                dt_used = dt_actual
            
            # Update solution
            U, S, V = U_new, S_new, V_new
            t += dt_used
            step_count += 1
            
            # Store history
            self.time_history.append(t)
            self.rank_history.append(S.shape[0])
            
            # Compute energy if available
            if hasattr(self.linear_solver, 'energy'):
                energy = self.linear_solver.energy(U, S, V)
                self.energy_history.append(energy)
            
            # Store solution
            if step_count % self.config.save_frequency == 0:
                solution = U @ S @ V.T
                self.solution_history.append(solution.copy())
            
            # Progress output
            if self.config.verbose and step_count % self.config.plot_frequency == 0:
                current_rank = S.shape[0]
                elapsed_time = time.time() - start_time
                print(f"  t = {t:.3f}, rank = {current_rank}, dt = {dt_used:.2e}, "
                      f"elapsed = {elapsed_time:.1f}s")
                
                # Optional: plot current solution
                if step_count % (self.config.plot_frequency * 5) == 0:
                    self._plot_current_solution(U @ S @ V.T, t)
        
        total_time = time.time() - start_time
        
        if self.config.verbose:
            print(f"DLR-FEM solution completed in {total_time:.2f}s")
            print(f"  Total steps: {step_count}")
            print(f"  Final rank: {S.shape[0]}")
            print(f"  Average rank: {np.mean(self.rank_history):.1f}")
        
        # Final solution
        final_solution = U @ S @ V.T
        
        return {
            'solution': final_solution,
            'U': U, 'S': S, 'V': V,
            'time_history': np.array(self.time_history),
            'rank_history': np.array(self.rank_history),
            'energy_history': np.array(self.energy_history) if self.energy_history else None,
            'solution_history': self.solution_history,
            'total_time': total_time,
            'total_steps': step_count,
            'mesh': self.mesh,
            'config': self.config
        }
    
    def _adaptive_step(self, U: np.ndarray, S: np.ndarray, V: np.ndarray,
                      dt: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float]:
        """Perform adaptive time step with error control."""
        # This could be enhanced with more sophisticated error estimation
        
        # For now, use the adaptive capability of individual solvers
        dt_used = dt
        
        # Try linear step first
        if hasattr(self.linear_solver, 'adaptive_step'):
            U_1, S_1, V_1, dt_linear = self.linear_solver.adaptive_step(
                U, S, V, dt/2, min_dt=self.config.min_dt
            )
            dt_used = min(dt_used, 2 * dt_linear)
        else:
            U_1, S_1, V_1 = self.linear_solver.step(U, S, V, dt/2)
        
        # Nonlinear step
        if hasattr(self.nonlinear_solver, 'adaptive_step'):
            U_2, S_2, V_2, dt_nonlinear = self.nonlinear_solver.adaptive_step(
                U_1, S_1, V_1, dt, min_dt=self.config.min_dt
            )
            dt_used = min(dt_used, dt_nonlinear)
        else:
            U_2, S_2, V_2 = self.nonlinear_solver.step(U_1, S_1, V_1, dt)
        
        # Second linear half-step
        U_final, S_final, V_final = self.linear_solver.step(U_2, S_2, V_2, dt_used/2)
        
        return U_final, S_final, V_final, dt_used
    
    def _plot_current_solution(self, solution: np.ndarray, t: float):
        """Plot current solution for monitoring."""
        plt.figure(figsize=(8, 6))
        plt.contourf(self.X, self.Y, solution, levels=20, cmap='coolwarm')
        plt.colorbar(label='u')
        plt.title(f'DLR-FEM Solution at t = {t:.3f}')
        plt.xlabel('x')
        plt.ylabel('y')
        plt.axis('equal')
        plt.tight_layout()
        plt.savefig('dlrfem_solution.png', dpi=300, bbox_inches='tight')
        print("Solution plot saved to dlrfem_solution.png")
        plt.close()
    
    def plot_results(self, results: Dict):
        """Plot comprehensive results."""
        fig, axes = plt.subplots(2, 3, figsize=(15, 10))
        
        # Final solution
        ax = axes[0, 0]
        im = ax.contourf(self.X, self.Y, results['solution'], levels=20, cmap='coolwarm')
        ax.set_title('Final Solution')
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.axis('equal')
        plt.colorbar(im, ax=ax)
        
        # Initial condition
        ax = axes[0, 1]
        im = ax.contourf(self.X, self.Y, self.u0, levels=20, cmap='coolwarm')
        ax.set_title('Initial Condition')
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.axis('equal')
        plt.colorbar(im, ax=ax)
        
        # Rank evolution
        ax = axes[0, 2]
        ax.plot(results['time_history'], results['rank_history'], 'b-', linewidth=2)
        ax.set_title('Rank Evolution')
        ax.set_xlabel('Time')
        ax.set_ylabel('Rank')
        ax.grid(True)
        
        # Energy evolution (if available)
        ax = axes[1, 0]
        if results['energy_history'] is not None:
            ax.plot(results['time_history'], results['energy_history'], 'r-', linewidth=2)
            ax.set_title('Energy Evolution')
            ax.set_xlabel('Time')
            ax.set_ylabel('Energy')
            ax.grid(True)
        else:
            ax.text(0.5, 0.5, 'Energy not computed', ha='center', va='center')
            ax.set_title('Energy Evolution')
        
        # Solution evolution (if stored)
        if len(self.solution_history) > 1:
            ax = axes[1, 1]
            # Plot solution at different times
            n_plots = min(4, len(self.solution_history))
            times_to_plot = np.linspace(0, len(self.solution_history)-1, n_plots, dtype=int)
            
            for i, idx in enumerate(times_to_plot):
                sol = self.solution_history[idx]
                ax.plot(sol[self.config.ny//2, :], label=f't={results["time_history"][idx*self.config.save_frequency]:.2f}')
            
            ax.set_title('Solution Cross-section (y=0.5)')
            ax.set_xlabel('x')
            ax.set_ylabel('u')
            ax.legend()
            ax.grid(True)
        
        # Statistics
        ax = axes[1, 2]
        stats_text = f"""DLR-FEM Statistics:
        
Final time: {results['time_history'][-1]:.3f}
Total steps: {results['total_steps']}
Computation time: {results['total_time']:.2f}s
        
Final rank: {results['rank_history'][-1]}
Average rank: {np.mean(results['rank_history']):.1f}
Max rank: {np.max(results['rank_history'])}
        
Grid: {self.config.nx}×{self.config.ny}
ε: {self.config.epsilon}
dt: {self.config.dt}"""
        
        ax.text(0.05, 0.95, stats_text, transform=ax.transAxes, 
                verticalalignment='top', fontfamily='monospace', fontsize=10)
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)
        ax.axis('off')
        
        plt.tight_layout()
        plt.savefig('dlrfem_solution.png', dpi=300, bbox_inches='tight')
        print("Solution plot saved to dlrfem_solution.png")
        plt.close()


def run_kiss_bubble_experiment(config: Optional[DLRFEMConfig] = None) -> Dict:
    """
    Run the kiss-bubble symmetry breaking experiment.
    
    Args:
        config: Optional configuration (uses default if None)
        
    Returns:
        Results dictionary
    """
    if config is None:
        config = DLRFEMConfig(
            nx=64, ny=64,
            T_final=0.5,
            dt=0.001,
            epsilon=0.02,
            max_rank=30,
            save_frequency=10
        )
    
    solver = DLRFEMSolver(config)
    
    # Run with kiss-bubble initial condition
    results = solver.solve(
        ic_type="kiss_bubble",
        centers=[(0.3, 0.5), (0.7, 0.5)],
        radius=0.15
    )
    
    # Plot results
    solver.plot_results(results)
    
    return results


if __name__ == "__main__":
    print("Running DLR-FEM example...")
    
    # Configure simulation
    config = DLRFEMConfig(
        nx=32, ny=32,  # Smaller grid for testing
        T_final=0.1,
        dt=0.005,
        epsilon=0.05,
        max_rank=20,
        save_frequency=5,
        plot_frequency=10,
        verbose=True
    )
    
    # Run kiss-bubble experiment
    results = run_kiss_bubble_experiment(config)
    
    print(f"\nSimulation completed successfully!")
    print(f"Final solution shape: {results['solution'].shape}")
    print(f"Final rank: {results['rank_history'][-1]}")
    print(f"Computation time: {results['total_time']:.2f}s")